package org.apache.flink.runtime.resourceconsumer.runner;


import java.util.Optional;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionException;
import java.util.concurrent.Executor;
import java.util.function.Supplier;

import javax.annotation.concurrent.GuardedBy;

import org.apache.flink.api.common.ConsumerID;
import org.apache.flink.api.common.ResourceStatus;
import org.apache.flink.api.common.time.Time;
import org.apache.flink.runtime.allocation.ResourceAllocation;
import org.apache.flink.runtime.allocation.ResourceAllocationStatistics;
import org.apache.flink.runtime.leaderelection.LeaderContender;
import org.apache.flink.runtime.leaderelection.LeaderElectionService;
import org.apache.flink.runtime.messages.Acknowledge;
import org.apache.flink.runtime.resourceconsumer.ConsumeCancellationFailedException;
import org.apache.flink.runtime.resourceconsumer.ResourceConsumerGateway;
import org.apache.flink.runtime.resourceconsumer.ResourceConsumerService;
import org.apache.flink.runtime.resourceconsumer.ResourceConsumerServiceProcess;
import org.apache.flink.runtime.resourceconsumer.factories.ResourceConsumerServiceProcessFactory;
import org.apache.flink.runtime.rpc.FatalErrorHandler;
import org.apache.flink.util.ExceptionUtils;
import org.apache.flink.util.ERException;
import org.apache.flink.util.Preconditions;
import org.apache.flink.util.concurrent.FutureUtils;
import org.apache.flink.util.function.ThrowingRunnable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;


/**
 * @类名: ResourceConsumerServiceLeadershipRunner
 * @说明: 消费者runner实现
 *
 * @author   leehom
 * @Date	 2023年11月1日 下午2:15:41
 * 修改记录：
 *
 * @see 	 
 */
/**
 * Leadership runner for the {@link ResourceConsumerServiceProcess}.
 * ResourceConsumerServiceProcess 领导权运行器，负责管理资源消费者领导权，选举和切换
 *
 * <p>The responsibility of this component is to manage the leadership of the {@link
 * ResourceConsumerServiceProcess}. This means that the runner will create an instance of the process when
 * it obtains the leadership. The process is stopped once the leadership is revoked.
 * 该组件的职责是管理ResourceConsumerServiceProcess的领导节点。
 *
 * <p>This component only accepts signals (consume result completion, initialization failure) as long as
 * it is running and as long as the signals are coming from the current leader process. This ensures
 * that only the current leader can affect this component.
 *
 * <p>All leadership operations are serialized. This means that granting the leadership has to
 * complete before the leadership can be revoked and vice versa.
 *
 * <p>The {@link #resultFuture} can be completed with the following values: * *
 *
 * <ul>
 *   <li>{@link ResourceConsumerRunnerResult} to signal an initialization failure of the {@link
 *       ResourceConsumerService} or the completion of a job
 *   <li>{@link Exception} to signal an unexpected failure
 * </ul>
 */ 
public class ResourceConsumerServiceLeadershipRunner implements ResourceConsumerRunner, LeaderContender {

    private static final Logger LOG =
            LoggerFactory.getLogger(ResourceConsumerServiceLeadershipRunner.class);

    private final Object lock = new Object();

    /** 资源消费者服务处理工厂*/
    private final ResourceConsumerServiceProcessFactory resourceConsumerServiceProcessFactory;
    /** 主节点选举服务*/
    private final LeaderElectionService leaderElectionService;

    private final FatalErrorHandler fatalErrorHandler;

    /** 异步关闭关闭结果*/
    private final CompletableFuture<Void> terminationFuture = new CompletableFuture<>();

    /** runner执行结果*/
    private final CompletableFuture<ResourceConsumerRunnerResult> resultFuture = new CompletableFuture<>();

    @GuardedBy("lock")
    private State state = State.RUNNING;
    /** */
    @GuardedBy("lock")
    private CompletableFuture<Void> sequentialOperation = FutureUtils.completedVoidFuture();
    /** */
    @GuardedBy("lock")
    private ResourceConsumerServiceProcess resourceConsumerServiceProcess =
            									ResourceConsumerServiceProcess.waitingForLeadership();
    /** 异步获取资源消费者Gateway*/
    @GuardedBy("lock")
    private CompletableFuture<ResourceConsumerGateway> resourceConsumerGatewayFuture = new CompletableFuture<>();

    @GuardedBy("lock")
    private boolean hasCurrentLeaderBeenCancelled = false;

    public ResourceConsumerServiceLeadershipRunner(
    		ResourceConsumerServiceProcessFactory resourceConsumerServiceProcessFactory,
            LeaderElectionService leaderElectionService,
            FatalErrorHandler fatalErrorHandler) {
        this.resourceConsumerServiceProcessFactory = resourceConsumerServiceProcessFactory;
        this.leaderElectionService = leaderElectionService;
        this.fatalErrorHandler = fatalErrorHandler;
    }

    /**
     * 关闭runner
     */
    @Override
    public CompletableFuture<Void> closeAsync() {
        final CompletableFuture<Void> processTerminationFuture;
        synchronized (lock) {
            if (state == State.STOPPED) {
                return terminationFuture;
            }
            state = State.STOPPED;
            LOG.debug("Terminating the leadership runner for consumer {}.", getConsumerID());
            // 
            resourceConsumerGatewayFuture.completeExceptionally(
                    new ERException(
                            "ResourceConsumerServiceLeadershipRunner is closed. "
                            + "Therefore, the corresponding ResourceConsumer will never acquire the leadership."));
            resultFuture.complete(
            		ResourceConsumerRunnerResult.forSuccess(
                            null
                            ));
            // 关闭服务处理
            processTerminationFuture = resourceConsumerServiceProcess.closeAsync();
        }
        // 
        final CompletableFuture<Void> serviceTerminationFuture =
                FutureUtils.runAfterwards(
                        processTerminationFuture,
                        () -> {
                            leaderElectionService.stop();
                        });
        //
        FutureUtils.forward(serviceTerminationFuture, terminationFuture);
        //
        terminationFuture.whenComplete(
                (unused, throwable) ->
                        LOG.debug("Leadership runner for consumer {} has been terminated.", getConsumerID()));
        return terminationFuture;
    }

    @Override
    public void start() throws Exception {
        LOG.debug("Start leadership runner for consumer {}.", getConsumerID());
        leaderElectionService.start(this);
    }

    // 
    @Override
    public CompletableFuture<ResourceConsumerGateway> getResourceConsumerGateway() {
        synchronized (lock) {
            return resourceConsumerGatewayFuture;
        }
    }

    @Override
    public CompletableFuture<ResourceConsumerRunnerResult> getResultFuture() {
        return resultFuture;
    }

    @Override
    public ConsumerID getConsumerID() {
        return resourceConsumerServiceProcessFactory.getConsumerId();
    }

    @Override
    public CompletableFuture<Acknowledge> cancel(Time timeout) {
        synchronized (lock) {
            hasCurrentLeaderBeenCancelled = true;
            return getResourceConsumerGateway()
                    .thenCompose(resourceConsumerGateway -> resourceConsumerGateway.cancel(timeout))
                    .exceptionally(
                            e -> {
                                throw new CompletionException(
                                        new ConsumeCancellationFailedException(
                                                "Cancellation failed.",
                                                ExceptionUtils.stripCompletionException(e)));
                            });
        }
    }

    @Override
    public boolean isInitialized() {
        synchronized (lock) {
            return resourceConsumerServiceProcess.isInitializedAndRunning();
        }
    }

    @Override
    public void grantLeadership(UUID leaderSessionID) {
    	// 启动资源消费者服务process
        runIfStateRunning(
                () -> startResourceConsumerServiceProcessAsync(leaderSessionID),
                "starting a new ResourceConsumerServiceProcess");
    }

    // 启动资源消费者服务process
    @GuardedBy("lock")
    private void startResourceConsumerServiceProcessAsync(UUID leaderSessionId) {
        sequentialOperation =
                sequentialOperation.thenRun(
                        () ->
                                runIfValidLeader(
                                        leaderSessionId,
                                        ThrowingRunnable.unchecked(
                                                () ->
                                                        verifyConsumeStatusAndCreateResourceConsumerServiceProcess(
                                                                leaderSessionId)),
                                        "verify resource allocating status and create ResourceConsumerServiceProcess"));
        handleAsyncOperationError(sequentialOperation, "Could not start the resource consumer.");
    }

    // 确认资源消费状态，构建资源消费者服务，故障恢复
    @GuardedBy("lock")
    private void verifyConsumeStatusAndCreateResourceConsumerServiceProcess(UUID leaderSessionId)
            throws ERException {
    	createNewResourceConsumerServiceProcess(leaderSessionId);
    }

    @GuardedBy("lock")
    private void createNewResourceConsumerServiceProcess(UUID leaderSessionId) throws ERException {
        Preconditions.checkState(resourceConsumerServiceProcess.closeAsync().isDone());
        LOG.info(
                "{} for consumer {} was granted leadership with leader id {}. Creating new {}.",
                getClass().getSimpleName(),
                getConsumerID(),
                leaderSessionId,
                ResourceConsumerServiceProcess.class.getSimpleName());
        //
        resourceConsumerServiceProcess = resourceConsumerServiceProcessFactory.create(leaderSessionId);
        //
        forwardIfValidLeader(
                leaderSessionId,
                resourceConsumerServiceProcess.getResourceConsumerGatewayFuture(),
                resourceConsumerGatewayFuture,
                "JobMasterGatewayFuture from JobMasterServiceProcess");
        //
        forwardResultFuture(leaderSessionId, resourceConsumerServiceProcess.getResultFuture());
        // 写入地址信息 
        confirmLeadership(leaderSessionId, resourceConsumerServiceProcess.getLeaderAddressFuture());
    }

    private void confirmLeadership(
            UUID leaderSessionId, CompletableFuture<String> leaderAddressFuture) {
        FutureUtils.assertNoException(
                leaderAddressFuture.thenAccept(
                        address -> {
                            synchronized (lock) {
                                if (isValidLeader(leaderSessionId)) {
                                    LOG.debug("Confirm leadership {}.", leaderSessionId);
                                    leaderElectionService.confirmLeadership(
                                            leaderSessionId, address);
                                } else {
                                    LOG.trace(
                                            "Ignore confirming leadership because the leader {} is no longer valid.",
                                            leaderSessionId);
                                }
                            }
                        }));
    }

    private void forwardResultFuture(
            UUID leaderSessionId, CompletableFuture<ResourceConsumerRunnerResult> resultFuture) {
        resultFuture.whenComplete(
                (resourceConsumerRunnerResult, throwable) -> {
                    synchronized (lock) {
                        if (isValidLeader(leaderSessionId)) {
                            onConsumeCompletion(resourceConsumerRunnerResult, throwable);
                        } else {
                            LOG.trace(
                                    "Ignore result future forwarding because the leader {} is no longer valid.",
                                    leaderSessionId);
                        }
                    }
                });
    }

    @GuardedBy("lock")
    private void onConsumeCompletion(
            ResourceConsumerRunnerResult resourceConsumerRunnerResult, Throwable throwable) {
        state = State.COMPLETED;

        LOG.debug("Completing the result for consumer {}.", getConsumerID());
        // 
        if (throwable != null) {
            resultFuture.completeExceptionally(throwable);
            resourceConsumerGatewayFuture.completeExceptionally(
                    new ERException(
                            "Could not retrieve ResourceConsumerGateway because the ResourceConsumer failed.",
                            throwable));
        } else {
            if (!resourceConsumerRunnerResult.isSuccess()) {
                resourceConsumerGatewayFuture.completeExceptionally(
                        new ERException(
                                "Could not retrieve JobMasterGateway because the JobMaster initialization failed.",
                                resourceConsumerRunnerResult.getInitializationFailure()));
            }

            resultFuture.complete(resourceConsumerRunnerResult);
        }
    }

    @Override
    public void revokeLeadership() {
        runIfStateRunning(
                this::stopResourceConsumerServiceProcessAsync,
                "revoke leadership from ResourceConsumerServiceProcess");
    }

    // 异步停止资源消费者服务处理
    @GuardedBy("lock")
    private void stopResourceConsumerServiceProcessAsync() {
        sequentialOperation =
                sequentialOperation.thenCompose(
                        ignored ->
                                callIfRunning(
                                                this::stopResourceConsumerServiceProcess,
                                                "stop leading JobMasterServiceProcess")
                                        .orElse(FutureUtils.completedVoidFuture()));

        handleAsyncOperationError(sequentialOperation, "Could not suspend the job manager.");
    }

    @GuardedBy("lock")
    private CompletableFuture<Void> stopResourceConsumerServiceProcess() {
        LOG.info(
                "{} for consumer {} was revoked leadership with leader id {}. Stopping current {}.",
                getClass().getSimpleName(),
                getConsumerID(),
                resourceConsumerServiceProcess.getLeaderSessionId(),
                ResourceConsumerServiceProcess.class.getSimpleName());

        resourceConsumerGatewayFuture.completeExceptionally(
                new ERException(
                        "Cannot obtain JobMasterGateway because the JobMaster lost leadership."));
        resourceConsumerGatewayFuture = new CompletableFuture<>();

        hasCurrentLeaderBeenCancelled = false;

        return resourceConsumerServiceProcess.closeAsync();
    }

    @Override
    public void handleError(Exception exception) {
        fatalErrorHandler.onFatalError(exception);
    }

    private void handleAsyncOperationError(CompletableFuture<Void> operation, String message) {
        operation.whenComplete(
                (unused, throwable) -> {
                    if (throwable != null) {
                        runIfStateRunning(
                                () ->
                                        handleServiceLeadershipRunnerError(
                                                new ERException(message, throwable)),
                                "handle ResourceConsumerServiceLeadershipRunner error");
                    }
                });
    }

    private void handleServiceLeadershipRunnerError(Throwable cause) {
        if (ExceptionUtils.isJvmFatalError(cause)) {
            fatalErrorHandler.onFatalError(cause);
        } else {
            resultFuture.completeExceptionally(cause);
        }
    }

    private void runIfStateRunning(Runnable action, String actionDescription) {
        synchronized (lock) {
            if (isRunning()) {
                action.run();
            } else {
                LOG.trace(
                        "Ignore '{}' because the leadership runner is no longer running.",
                        actionDescription);
            }
        }
    }

    private <T> Optional<T> callIfRunning(
            Supplier<? extends T> supplier, String supplierDescription) {
        synchronized (lock) {
            if (isRunning()) {
                return Optional.of(supplier.get());
            } else {
                LOG.trace(
                        "Ignore '{}' because the leadership runner is no longer running.", supplierDescription);
                return Optional.empty();
            }
        }
    }

    @GuardedBy("lock")
    private boolean isRunning() {
        return state == State.RUNNING;
    }

    private void runIfValidLeader(
            UUID expectedLeaderId, Runnable action, String actionDescription) {
        synchronized (lock) {
            if (isValidLeader(expectedLeaderId)) {
                action.run();
            } else {
                LOG.trace(
                        "Ignore leader action '{}' because the leadership runner is no longer the valid leader for {}.",
                        actionDescription,
                        expectedLeaderId);
            }
        }
    }

    @GuardedBy("lock")
    private boolean isValidLeader(UUID expectedLeaderId) {
        return isRunning() && leaderElectionService.hasLeadership(expectedLeaderId);
    }

    private <T> void forwardIfValidLeader(
            UUID expectedLeaderId,
            CompletableFuture<? extends T> source,
            CompletableFuture<T> target,
            String forwardDescription) {
        source.whenComplete(
                (t, throwable) -> {
                    synchronized (lock) {
                        if (isValidLeader(expectedLeaderId)) {
                            if (throwable != null) {
                                target.completeExceptionally(throwable);
                            } else {
                                target.complete(t);
                            }
                        } else {
                            LOG.trace(
                                    "Ignore forwarding '{}' because the leadership runner is no longer the valid leader for {}.",
                                    forwardDescription,
                                    expectedLeaderId);
                        }
                    }
                });
    }

    enum State {
        RUNNING,
        STOPPED,
        COMPLETED,
    }

	@Override
	public CompletableFuture<Void> localCleanupAsync(ConsumerID jobId, Executor cleanupExecutor) {
		return this.closeAsync();
	}
}
